Dostrojenie modelu Gemma za pomocą funkcji Hugging Face Transformers i QloRA

Z tego przewodnika dowiesz się, jak dostosować ustawienia Gemma do niestandardowego zbioru danych tekst-SQL za pomocą Transformerów Hugging Face i TRL. Dowiesz się:

  • Co to jest kwantyzowana adaptacja niskiego rzędu (QLoRA)
  • Konfigurowanie środowiska programistycznego
  • Tworzenie i przygotowywanie zbioru danych do dostrajania
  • Dostrojenie Gemma za pomocą TRL i SFTTrainer
  • Testowanie wnioskowania modelu i generowanie zapytań SQL

Co to jest kwantyzacja adaptacji niskiego rzędu

Ten przewodnik pokazuje, jak korzystać z zaadaptowanej kwantyzowanej niskorzędowej adaptacji (QLoRA), która stała się popularną metodą efektywnego dopracowywania modeli LLM, ponieważ zmniejsza wymagania dotyczące zasobów obliczeniowych przy jednoczesnym zachowaniu wysokiej wydajności. W QloRA wytrenowany model jest kwantyzowany do 4-bitów, a wagi są blokowane. Następnie dołączane są warstwy adaptera, które można trenować (LoRA), i tylko te warstwy są trenowane. Następnie wagi adaptera można scalić z modelem podstawowym lub zachować jako osobny adapter.

Konfigurowanie środowiska programistycznego

Pierwszym krokiem jest zainstalowanie bibliotek Hugging Face, w tym TRL, oraz zbiorów danych, aby dostosować model otwarty, w tym różne techniki RLHF i wyrównywania.

# Install Pytorch & other libraries
%pip install "torch>=2.4.0" tensorboard

# Install Gemma release branch from Hugging Face
%pip install "transformers>=4.51.3"

# Install Hugging Face libraries
%pip install  --upgrade \
  "datasets==3.3.2" \
  "accelerate==1.4.0" \
  "evaluate==0.4.3" \
  "bitsandbytes==0.45.3" \
  "trl==0.15.2" \
  "peft==0.14.0" \
  protobuf \
  sentencepiece

# COMMENT IN: if you are running on a GPU that supports BF16 data type and flash attn, such as NVIDIA L4 or NVIDIA A100
#% pip install flash-attn

Uwaga: jeśli używasz karty graficznej z architekturą Ampere (np. NVIDIA L4) lub nowszą, możesz korzystać z funkcji Flash attention. Flash Attention to metoda, która znacznie przyspiesza obliczenia i zmniejsza wykorzystanie pamięci z kwadratowego na liniowe w długości sekwencji, co prowadzi do przyspieszenia treningu nawet 3-krotnie. Więcej informacji znajdziesz w artykule FlashAttention.

Zanim zaczniesz szkolenie, musisz zaakceptować warunki korzystania z usługi Gemma. Umowę licencyjną możesz zaakceptować w Hugging Face, klikając przycisk Zgadzam się i uzyskaj dostęp do repozytorium na stronie modelu pod adresem: http://huggingface.co/google/gemma-3-1b-pt

Po zaakceptowaniu licencji musisz mieć ważny token Hugging Face, aby uzyskać dostęp do modelu. Jeśli używasz Google Colab, możesz bezpiecznie korzystać z tokenu Hugging Face, korzystając z sekretów Colab. Możesz też ustawić token bezpośrednio w metodzie login. Podczas przesyłania modelu do Huba podczas trenowania upewnij się, że token ma też uprawnienia do zapisu.

from google.colab import userdata
from huggingface_hub import login

# Login into Hugging Face Hub
hf_token = userdata.get('HF_TOKEN') # If you are running inside a Google Colab
login(hf_token)

Tworzenie i przygotowywanie zbioru danych do dostrajania

Podczas dopracowywania LLM ważne jest, aby znać przypadek użycia i zadania, które chcesz wykonać. Pomaga to utworzyć zbiór danych, który posłuży do dokładnego dostosowania modelu. Jeśli nie masz jeszcze zdefiniowanego przypadku użycia, możesz wrócić do etapu projektowania.

W tym przewodniku omówiono następujący przypadek użycia:

  • Dostroić model języka naturalnego do modelu SQL, aby umożliwić płynną integrację z narzędziem do analizy danych. Celem jest znaczne skrócenie czasu i zmniejszenie zapotrzebowania na wiedzę specjalistyczną w przypadku generowania zapytań SQL, co umożliwi nawet użytkownikom bez wiedzy technicznej wydobywanie z danych przydatnych informacji.

Konwersja tekstu na SQL może być przydatna do dokładnego dostrajania LLM, ponieważ jest to złożone zadanie, które wymaga dużej (wewnętrznej) wiedzy o danych i języku SQL.

Gdy uznasz, że dostosowanie jest odpowiednim rozwiązaniem, musisz mieć zbiór danych, który będziesz dostosowywać. Zestaw danych powinien zawierać różnorodne demonstracje zadań, które chcesz rozwiązać. Taki zbiór danych możesz utworzyć na kilka sposobów:

  • Korzystanie z dostępnych zbiorów danych typu open source, np. Spider
  • Korzystanie ze zbiorów danych syntetycznych utworzonych przez modele LLM, np. Alpaca.
  • Korzystanie ze zbiorów danych utworzonych przez ludzi, takich jak Dolly.
  • Korzystanie z kombinacji metod, takich jak Orca

Każda z tych metod ma swoje zalety i wady oraz zależy od budżetu, czasu i wymagań dotyczących jakości. Na przykład korzystanie z dotychczasowego zbioru danych jest najprostsze, ale może nie być dostosowane do Twojego konkretnego przypadku użycia. Z kolei korzystanie z usług ekspertów w danej dziedzinie może być najbardziej dokładne, ale może zająć dużo czasu i być kosztowne. Można też połączyć kilka metod, aby utworzyć zbiór danych instrukcji, jak pokazano w artykule Orca: Progressive Learning from Complex Explanation Traces of GPT-4 (w języku angielskim).

Ten przewodnik korzysta z już istniejącego zbioru danych (philschmid/gretel-synthetic-text-to-sql), wysokiej jakości syntetycznego zbioru danych Text-to-SQL, który zawiera instrukcje w języku naturalnym, definicje schematów, rozumowanie i odpowiednie zapytanie SQL.

Hugging Face TRL obsługuje automatyczne tworzenie szablonów formatów zbiorów danych rozmów. Oznacza to, że wystarczy przekonwertować zbiór danych na odpowiednie obiekty JSON, a trl zajmie się tworzeniem szablonów i przekształcaniem ich w odpowiednie formaty.

{"messages": [{"role": "system", "content": "You are..."}, {"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]}
{"messages": [{"role": "system", "content": "You are..."}, {"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]}
{"messages": [{"role": "system", "content": "You are..."}, {"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]}

Pakiet philschmid/gretel-synthetic-text-to-sql zawiera ponad 100 tys. próbek. Aby zachować mały rozmiar przewodnika, zmniejszyliśmy próbkowanie,aby zawierał tylko 10 tys. próbek.

Teraz możesz użyć biblioteki Hugging Face Datasets, aby załadować zbiór danych i utworzyć szablon promptu, który łączy instrukcje w naturalnym języku z definicją schematu i dodaje komunikat systemowy dla asystenta.

from datasets import load_dataset

# System message for the assistant
system_message = """You are a text to SQL query translator. Users will ask you questions in English and you will generate a SQL query based on the provided SCHEMA."""

# User prompt that combines the user query and the schema
user_prompt = """Given the <USER_QUERY> and the <SCHEMA>, generate the corresponding SQL command to retrieve the desired data, considering the query's syntax, semantics, and schema constraints.

<SCHEMA>
{context}
</SCHEMA>

<USER_QUERY>
{question}
</USER_QUERY>
"""
def create_conversation(sample):
  return {
    "messages": [
      # {"role": "system", "content": system_message},
      {"role": "user", "content": user_prompt.format(question=sample["sql_prompt"], context=sample["sql_context"])},
      {"role": "assistant", "content": sample["sql"]}
    ]
  }

# Load dataset from the hub
dataset = load_dataset("philschmid/gretel-synthetic-text-to-sql", split="train")
dataset = dataset.shuffle().select(range(12500))

# Convert dataset to OAI messages
dataset = dataset.map(create_conversation, remove_columns=dataset.features,batched=False)
# split dataset into 10,000 training samples and 2,500 test samples
dataset = dataset.train_test_split(test_size=2500/12500)

# Print formatted user prompt
print(dataset["train"][345]["messages"][1]["content"])

Dostrojenie Gemma za pomocą TRL i SFTTrainer

Możesz teraz dostosować model. Narzędzie Hugging Face TRL SFTTrainer ułatwia nadzorowanie dopracowywania otwartych modeli LLM. SFTTrainer jest podklasą Trainer z biblioteki transformers i obsługuje te same funkcje, w tym rejestrowanie, ocenę i punkty kontrolne, ale zawiera też dodatkowe funkcje, takie jak:

  • Formatowanie zbioru danych, w tym formaty konwersacyjne i instrukcyjne
  • Trenowanie tylko na podstawie ukończonych zadań, z pominięciem promptów
  • Pakowanie zbiorów danych w celu efektywniejszego trenowania
  • obsługa dostrajania konkretnych parametrów (PEFT), w tym QloRA;
  • Przygotowanie modelu i tokenizera do dokładnego dostrajania konwersacji (np. dodawanie tokenów specjalnych)

Poniższy kod wczytuje model Gemma i tokenizer z Hugging Face oraz inicjuje konfigurację kwantyzacji.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModelForImageTextToText, BitsAndBytesConfig

# Hugging Face model id
model_id = "google/gemma-3-1b-pt" # or `google/gemma-3-4b-pt`, `google/gemma-3-12b-pt`, `google/gemma-3-27b-pt`

# Select model class based on id
if model_id == "google/gemma-3-1b-pt":
    model_class = AutoModelForCausalLM
else:
    model_class = AutoModelForImageTextToText

# Check if GPU benefits from bfloat16
if torch.cuda.get_device_capability()[0] >= 8:
    torch_dtype = torch.bfloat16
else:
    torch_dtype = torch.float16

# Define model init arguments
model_kwargs = dict(
    attn_implementation="eager", # Use "flash_attention_2" when running on Ampere or newer GPU
    torch_dtype=torch_dtype, # What torch dtype to use, defaults to auto
    device_map="auto", # Let torch decide how to load the model
)

# BitsAndBytesConfig: Enables 4-bit quantization to reduce model size/memory usage
model_kwargs["quantization_config"] = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type='nf4',
    bnb_4bit_compute_dtype=model_kwargs['torch_dtype'],
    bnb_4bit_quant_storage=model_kwargs['torch_dtype'],
)

# Load model and tokenizer
model = model_class.from_pretrained(model_id, **model_kwargs)
tokenizer = AutoTokenizer.from_pretrained("google/gemma-3-1b-it") # Load the Instruction Tokenizer to use the official Gemma template

SFTTrainer obsługuje natywną integrację z peft, co ułatwia skuteczne dostrajanie LLM za pomocą QLoRA. Wystarczy, że utworzysz LoraConfig i przekażesz go trenerowi.

from peft import LoraConfig

peft_config = LoraConfig(
    lora_alpha=16,
    lora_dropout=0.05,
    r=16,
    bias="none",
    target_modules="all-linear",
    task_type="CAUSAL_LM",
    modules_to_save=["lm_head", "embed_tokens"] # make sure to save the lm_head and embed_tokens as you train the special tokens
)

Zanim rozpoczniesz szkolenie, musisz zdefiniować hiperparametr, którego chcesz użyć w instancji SFTConfig.

from trl import SFTConfig

args = SFTConfig(
    output_dir="gemma-text-to-sql",         # directory to save and repository id
    max_seq_length=512,                     # max sequence length for model and packing of the dataset
    packing=True,                           # Groups multiple samples in the dataset into a single sequence
    num_train_epochs=3,                     # number of training epochs
    per_device_train_batch_size=1,          # batch size per device during training
    gradient_accumulation_steps=4,          # number of steps before performing a backward/update pass
    gradient_checkpointing=True,            # use gradient checkpointing to save memory
    optim="adamw_torch_fused",              # use fused adamw optimizer
    logging_steps=10,                       # log every 10 steps
    save_strategy="epoch",                  # save checkpoint every epoch
    learning_rate=2e-4,                     # learning rate, based on QLoRA paper
    fp16=True if torch_dtype == torch.float16 else False,   # use float16 precision
    bf16=True if torch_dtype == torch.bfloat16 else False,   # use bfloat16 precision
    max_grad_norm=0.3,                      # max gradient norm based on QLoRA paper
    warmup_ratio=0.03,                      # warmup ratio based on QLoRA paper
    lr_scheduler_type="constant",           # use constant learning rate scheduler
    push_to_hub=True,                       # push model to hub
    report_to="tensorboard",                # report metrics to tensorboard
    dataset_kwargs={
        "add_special_tokens": False, # We template with special tokens
        "append_concat_token": True, # Add EOS token as separator token between examples
    }
)

Masz już wszystkie elementy, których potrzebujesz do utworzenia SFTTrainer, aby rozpocząć trenowanie modelu.

from trl import SFTTrainer

# Create Trainer object
trainer = SFTTrainer(
    model=model,
    args=args,
    train_dataset=dataset["train"],
    peft_config=peft_config,
    processing_class=tokenizer
)

Rozpocznij trenowanie, wywołując metodę train().

# Start training, the model will be automatically saved to the Hub and the output directory
trainer.train()

# Save the final model again to the Hugging Face Hub
trainer.save_model()

Zanim przetestujesz model, musisz zwolnić pamięć.

# free the memory again
del model
del trainer
torch.cuda.empty_cache()

Korzystając z QLoRA, trenujesz tylko adaptery, a nie cały model. Oznacza to, że podczas zapisywania modelu podczas trenowania zapisujesz tylko wagi adaptera, a nie cały model. Jeśli chcesz zapisać pełny model, co ułatwia jego używanie z usługami warstwy serwowania, takimi jak vLLM czy TGI, możesz scalić wagi adaptera z wagami modelu za pomocą metody merge_and_unload, a następnie zapisać model za pomocą metody save_pretrained. Dzięki temu zapisywany jest model domyślny, który może być używany do wnioskowania.

from peft import PeftModel

# Load Model base model
model = model_class.from_pretrained(model_id, low_cpu_mem_usage=True)

# Merge LoRA and base model and save
peft_model = PeftModel.from_pretrained(model, args.output_dir)
merged_model = peft_model.merge_and_unload()
merged_model.save_pretrained("merged_model", safe_serialization=True, max_shard_size="2GB")

processor = AutoTokenizer.from_pretrained(args.output_dir)
processor.save_pretrained("merged_model")

Testowanie wnioskowania modelu i generowanie zapytań SQL

Po zakończeniu trenowania należy ocenić i przetestować model. Możesz wczytywać różne próbki ze zbioru danych testowego i ocenić model na ich podstawie.

import torch
from transformers import pipeline

model_id = "gemma-text-to-sql"

# Load Model with PEFT adapter
model = model_class.from_pretrained(
  model_id,
  device_map="auto",
  torch_dtype=torch_dtype,
  attn_implementation="eager",
)
tokenizer = AutoTokenizer.from_pretrained(model_id)

Wczytajmy losową próbę z zbioru danych testowego i wygenerujmy polecenie SQL.

from random import randint
import re

# Load the model and tokenizer into the pipeline
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

# Load a random sample from the test dataset
rand_idx = randint(0, len(dataset["test"]))
test_sample = dataset["test"][rand_idx]

# Convert as test example into a prompt with the Gemma template
stop_token_ids = [tokenizer.eos_token_id, tokenizer.convert_tokens_to_ids("<end_of_turn>")]
prompt = pipe.tokenizer.apply_chat_template(test_sample["messages"][:2], tokenize=False, add_generation_prompt=True)

# Generate our SQL query.
outputs = pipe(prompt, max_new_tokens=256, do_sample=False, temperature=0.1, top_k=50, top_p=0.1, eos_token_id=stop_token_ids, disable_compile=True)

# Extract the user query and original answer
print(f"Context:\n", re.search(r'<SCHEMA>\n(.*?)\n</SCHEMA>', test_sample['messages'][0]['content'], re.DOTALL).group(1).strip())
print(f"Query:\n", re.search(r'<USER_QUERY>\n(.*?)\n</USER_QUERY>', test_sample['messages'][0]['content'], re.DOTALL).group(1).strip())
print(f"Original Answer:\n{test_sample['messages'][1]['content']}")
print(f"Generated Answer:\n{outputs[0]['generated_text'][len(prompt):].strip()}")

Podsumowanie i kolejne kroki

W tym samouczku pokazaliśmy, jak dostrajać model Gemma za pomocą TRL i QLoRA. Zapoznaj się z tymi dokumentami: